#!/usr/bin/perl

# ------------------------------------------------------------------------------
# --- Creates an individual PDB file for all partner combos of each multichain -
# --- PDB file. Also creates a pdb file for each individual chain. -------------
# --- Original by Andrea -------------------------------------------------------
# --- Re-written by C. Bergey; documentation added by A. Watkins and C. Bergey -
# ------------------------------------------------------------------------------

use strict;
use warnings;

my $mincode;
my $argc = scalar (@ARGV);

if ($argc < 1 || $argc == 2 || $argc > 3) {
	print "Usage error. Either provide PDB code (four-letter) or PDB code and both chain codes.\n";
	print "Example: perl pdb_partner_single.pl 1YCR \n";
	print "Example: perl pdb_partner_single.pl 1YCR A B\n";
	exit 1;
}

my $pdb = $ARGV[0];
my $chain_1;
my $chain_2;
if (scalar (@ARGV) == 3) {
	$chain_1 = $ARGV[1];
	$chain_2 = $ARGV[2];

}

# Generate filename from PDB code; open it.
my $pdb_file = "pdb_files/$pdb\.pdb";
open(PDB, "<$pdb_file")
  or die "ERROR: Could not open PDB file ($pdb).\n";
open NORMALOUT, ">pdb_files/$pdb\.clean\.pdb";

my $model_count = 0;
foreach (<PDB>) {
	if (/^MODEL/) {
		# this file needs to be split
       		$model_count++;
		system ("mkdir pdb_files/$pdb" . "_models");
		open MDLOUT, ">pdb_files/$pdb" . "_models/$pdb" . "_m$model_count.pdb";
	}
	if ($model_count != 0) {
		if (/^TER/ or  /^ATOM/) {
			print MDLOUT $_;
		}
	}
	else {
		if (/^TER/ or  /^ATOM/) {
			print NORMALOUT $_;
		}

	}
	if (/^ENDMDL/) {
		close MDLOUT;
	}
}

close NORMALOUT;

my $pdb_final;

if ($model_count != 0) {
	# we need to pre-process FOR MODELS
	for (my $i = 1; $i <= $model_count; $i++) {
		my $ros_cmd = "/scratch/cmb433/pdb_interact/rosetta3.4/rosetta_source/bin/relax.linuxgccrelease "
		         . "-database /scratch/cmb433/pdb_interact/rosetta3.4/rosetta_database/ "
		         . "-relax:constrain_relax_to_start_coords -in:file:fullatom -nstruct 50 "
		         . "-in:file:s pdb_files/$pdb" . "_models/$pdb" . "_m$i.pdb";
		system ($ros_cmd);
	}
	
	my @files = <$pdb*.pdb>;
	my %scores;

	foreach my $file (@files) {
	        #print $file . "\n";
	        open FILE, $file;
	        while (<FILE>) {
	                if (/^pose/) {
	                        my @lineobjs = split;
				$scores{$file} = $lineobjs[20];
                	}
        	}

        	close (FILE);
	}

	foreach my $file (sort {$scores{$a} <=> $scores{$b}} keys %scores) {
		$pdb_final = $file;
		last;
	}

} else {

	# we need to pre-process
	
	my $ros_cmd = "/scratch/cmb433/pdb_interact/rosetta3.4/rosetta_source/bin/relax.linuxgccrelease "
	         . "-database /scratch/cmb433/pdb_interact/rosetta3.4/rosetta_database/ "
	         . "-relax:constrain_relax_to_start_coords -in:file:fullatom -nstruct 50 "
	         . "-in:file:s pdb_files/$pdb.clean.pdb";
	system ($ros_cmd);
	
	my @files = <$pdb.clean*.pdb>;
	my %scores;

	foreach my $file (@files) {
	        #print $file . "\n";
	        open FILE, $file;
	        while (<FILE>) {
	                if (/^pose/) {
	                        my @lineobjs = split;
				$scores{$file} = $lineobjs[20];
                	}
        	}

        	close (FILE);
	}

	foreach my $file (sort {$scores{$a} <=> $scores{$b}} keys %scores) {
		$pdb_final = $file;
		last;
	}
}

print $pdb_final . "\n";

# Store number of chains, and their IDs in a hash
my $nchain = 0;
my %chain_ids;
       
print STDERR "Processing PDB file $pdb_final...\n";
       
if ($argc == 1) {

	# we need to grab the chain letters

	open PDB, $pdb_final;	

	foreach (<PDB>) {
        
		if (/^ATOM/) {
			my $letter = substr($_,21,1);
			$chain_ids{$letter} = 1;
		}

		$nchain++ if (/^TER/);

	}

	close PDB;

}
else {
	$chain_ids{$chain_1} = 1;
	$chain_ids{$chain_2} = 1;
	$nchain = 2;
}


# Create, per PDB file, each of the possible n(n-1)/2 two-chain files
       
# Make and sort @all_chains, so we don't look for B_A instead of A_B
my @all_chains = sort(keys %chain_ids);
     
if (!-e "pdb_partner_files") {
	system ("mkdir ./pdb_partner_files");
}

# For all possible interchain partners...
for (my $chain_index_1 = 0; $chain_index_1 < scalar (@all_chains); $chain_index_1++) {
	for (my $chain_index_2 = $chain_index_1 + 1; $chain_index_2 < scalar (@all_chains); $chain_index_2++) {

		my $partner_1 = $all_chains[$chain_index_1];
		my $partner_2 = $all_chains[$chain_index_2];

		# Deduce output PDB filename
		my $out_file = "pdb_partner_files/" . $pdb;
		$out_file .= "_" . $all_chains[$chain_index_1];
		$out_file .= "_" . $all_chains[$chain_index_2];
		$out_file .= ".pdb";
		
		# Skip writing this PDB partner file if it already exists
		if (-e $out_file) {
			next;
		}
		
		# Create file and write to it
		open OUT_PDB, ">$out_file";

		# Print coordinates for first partner to file
		# Look for lines that start ATOM and contain the partner_1 chain ID
		# or partner_2 chain ID.
		
		my @partner_1_lines;
		my @partner_2_lines;
		
		open(PDB, "<$pdb_final")
		  or die "ERROR: Could not open PDB file ($pdb_final).\n";
 
		foreach (<PDB>) {
			my @info = split;
			if ($info[0] eq "ATOM") {
				if ($info[4] eq $partner_1) {
					push (@partner_1_lines, $_);
				} elsif ($info[4] eq $partner_2) {
					push (@partner_2_lines, $_);
				}
			}
		}

		close PDB;

		# Print these lines to the output file
		# Print TER at end of each chain
		print OUT_PDB @partner_1_lines;
		print OUT_PDB "TER \n";
		print OUT_PDB @partner_2_lines;
		print OUT_PDB "TER \n";
		close OUT_PDB;
	}
}
        # Finally, print coordinates of chains in individual files, one per chain.
       
        # So if PDB file 1A0H has chains A, B, D, and E. The files created above will be:
        #       1A0H_A_B.pdb            1A0H_B_D.pdb
        #       1A0H_A_D.pdb            1A0H_B_E.pdb
        #       1A0H_A_E.pdb            1A0H_D_E.pdb
       
        # This step creates the following individual files:
        #       1A0H_A.pdb
        #       1A0H_B.pdb
        #       1A0H_D.pdb
        #       1A0H_E.pdb
       
foreach my $chain (@all_chains) {
	my $ind_filename = "pdb_partner_files/" . $pdb . "_" . $chain . ".pdb";
	if (-e $ind_filename) {
		next;
	}

	open IND_FILE, ">$ind_filename"
	  or die "ERROR: Could not open individual PDB chain file ($ind_filename).\n";
               
	open(PDB, "<$pdb_final")
		or die "ERROR: Could not open PDB file ($pdb).\n";
     
	# Loop through original PDB, as before, looking for this chain's ATOM lines
	foreach (<PDB>) {
		my @info = split;
		if ($info[0] eq "ATOM") {
			if ($info[4] eq $chain) {
				print IND_FILE $_;
			}
		}
	}
	close PDB;

	print IND_FILE "TER \n";
	close IND_FILE;
}

exit 0;
